import os
import shutil
import psutil
import subprocess
import time
import signal
import errno
import numpy as np

# 获取CPU逻辑核心数
CPU_COUNT = int(psutil.cpu_count(logical=False) * 0.7)
t_wait = 20
limit_cpu_ave = 70
limit_cpu_max = 100
limit_mem_ave = 90
limit_mem_max = 100


# _process_ids = []
# process_names = ['cmd', 'MATLAB']
# 获取当前脚本文件的所在目录
dir_current_file = os.path.dirname(os.path.abspath(__file__))
dir_parent_1 = os.path.dirname(os.path.dirname(dir_current_file))
second_track_b_dir = os.path.join(dir_parent_1, "第二赛道-B卷")
scenarios_dir = os.path.join(second_track_b_dir, "Scenarios")
copy_path = os.path.join(dir_parent_1, "copy")


def is_file_locked(filepath):
    if not os.path.exists(filepath):
        return False
    newpath = filepath + ".locktest"
    try:
        os.rename(filepath, newpath)
        try:
            os.rename(newpath, filepath)
        except Exception as e:
            print(f"警告：检测占用时，{newpath} 改回 {filepath} 失败，请手动恢复！错误信息：{e}")
            return True
        return False
    except OSError as e:
        if e.errno in (errno.EACCES, errno.EPERM, 13):
            return True
        return False
    

def get_env_name():
    """
    获取用户指定的conda环境名称
    如果用户未输入，则t_wait秒后使用默认环境名称
    
    Returns:
        str: conda环境名称
    """
    computer_name = os.environ.get('COMPUTERNAME', 'Unknown')

    if computer_name == 'DESKTOP-184P33J':
        default_env = 'onsite_unstructured_25_py11'
    elif computer_name == 'DESKTOP-7MVV87O':
        default_env = 'yls_onsite'
        # default_env = 'onsite_unstructured_python311'
    print(f"请输入conda环境名称 (当前电脑为{computer_name}→默认: {default_env}，{t_wait}秒后自动使用默认值): ", end="", flush=True)
    user_input = ""
    
    # 等待用户输入，最多等待6秒
    for _ in range(100):  # 每0.1秒检查一次，总共6秒
        if os.name == 'nt':  # Windows
            import msvcrt
            if msvcrt.kbhit():
                user_input += msvcrt.getche().decode('utf-8')
                if user_input.endswith('\r'):  # 回车键
                    user_input = user_input[:-1]
                    print()  # 换行
                    break
        else:  # Unix/Linux
            import select
            import sys
            if select.select([sys.stdin], [], [], 0.1)[0]:
                user_input = sys.stdin.readline().strip()
                break
        time.sleep(0.1)
    
    if user_input.strip():
        return user_input.strip()
    print(f"\n使用默认环境: {default_env}")
    return default_env


def check_resource_usage():
    cpu_usage_queue = []
    mem_usage_queue = []

    # 检查资源占用情况
    for _ in range(20):
        cpu_usage_queue.append(psutil.cpu_percent())
        mem_usage_queue.append(psutil.virtual_memory().percent)
        time.sleep(0.1)

    # 计算平均CPU和内存使用率以及最高CPU和内存使用率
    avg_cpu_usage = np.mean(cpu_usage_queue[2:])
    max_cpu_usage = max(cpu_usage_queue[2:])
    avg_mem_usage = np.mean(mem_usage_queue[2:])
    max_mem_usage = max(mem_usage_queue[2:])

    if avg_cpu_usage > limit_cpu_ave or max_cpu_usage > limit_cpu_max or avg_mem_usage > limit_mem_ave or max_mem_usage > limit_mem_max:
        print(f"Ave CPU: {avg_cpu_usage:.1f}%,  Max CPU: {max_cpu_usage:.1f}%,  Ave ROM: {avg_mem_usage:.1f}%,  Max ROM: {max_mem_usage:.1f}%, Resource usage is too high")
        return True
    return False
 

# 更新当前运行的进程数
def update_running_processes():
    # 通过检查cmd.exe进程中包含"planner"的进程来估计当前运行的planner进程数
    running_count = 0
    for proc in psutil.process_iter(['name', 'cmdline']):
        try:
            if proc.info['name'] == 'cmd.exe' and proc.info['cmdline'] and any('planner' in cmd for cmd in proc.info['cmdline'] if cmd):
                running_count += 1
        except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
            pass
    
    return running_count



def delete_file_dir(folder_path, wait_print=True, default_delete=True):
    """
    删除指定文件夹中的所有内容或者删除指定的文件，
    如果wait_print为True，则等待用户确认，用户输入回车则立即删除，用户没有输入则t_wait秒后删除；
    如果default_delete为True，则默认删除，如果default_delete为False，则不删除；
    如果wait_print为False，则立即删除；

    Args:
        folder_path: 要清空的文件夹路径
        wait_print: 是否等待用户确认
        default_delete: 是否默认删除
    """
    try:
        if os.path.exists(folder_path):
            if wait_print:
                if default_delete:
                    print(f"是否删除 {folder_path} 的内容? (按回车确认，或等待{t_wait}秒自动删除，输入n取消)")
                else:
                    print(f"是否删除 {folder_path} 的内容? (输入y确认删除，或等待{t_wait}秒自动取消，按回车取消)")
                
                # 设置超时时间
                timeout = t_wait
                start_time = time.time()
                user_input = None
                
                # 等待用户输入，最多等待t_wait秒
                while time.time() - start_time < timeout:
                    if os.name == 'nt':  # Windows系统
                        import msvcrt
                        if msvcrt.kbhit():
                            user_input = msvcrt.getch().decode('utf-8').lower()
                            break
                    else:  # Unix/Linux系统
                        import select
                        import sys
                        if select.select([sys.stdin], [], [], 0)[0]:
                            user_input = sys.stdin.readline().strip().lower()
                            break
                    time.sleep(0.1)
                
                # 处理用户输入
                if user_input == 'n' and default_delete:
                    print(f"用户取消删除 {folder_path} 的内容")
                    return
                elif user_input == 'y' and not default_delete:
                    print(f"用户确认删除 {folder_path} 的内容")
                else:
                    # 超时或按回车，根据default_delete决定行为
                    if not default_delete:
                        print(f"未收到确认删除指令，取消删除 {folder_path} 的内容\n")
                        return
                    print(f"开始删除 {folder_path} 的内容...")
            
            # 执行删除操作
            if os.path.isdir(folder_path):
                for item in os.listdir(folder_path):
                    item_path = os.path.join(folder_path, item)
                    if os.path.isfile(item_path):
                        os.remove(item_path)
                    elif os.path.isdir(item_path):
                        shutil.rmtree(item_path)
                print(f"已清空 {folder_path} 文件夹")
            elif os.path.isfile(folder_path):
                os.remove(folder_path)
                print(f"已删除文件 {folder_path}")
        else:
            print(f"路径 {folder_path} 不存在，将创建此文件夹")
            # 仅在路径应为文件夹时创建
            if not os.path.splitext(folder_path)[1]:  # 没有文件扩展名，认为是文件夹
                os.makedirs(folder_path, exist_ok=True)
    except Exception as e:
        print(f"处理 {folder_path} 时出错: {e}")


def is_scenario_completed(directory_path):
    """
    检查目录中的main_print.log文件是否包含完成标记
    
    Args:
        directory_path: 要检查的目录路径
        
    Returns:
        bool: 如果找到完成标记返回True，否则返回False
    """
    log_file_path = os.path.join(directory_path, "main_print.log")
    if not os.path.exists(log_file_path):
        return False
        
    try:
        # 从文件末尾读取最后几行
        with open(log_file_path, 'r', encoding='utf-8', errors='ignore') as f:
            # 从文件末尾读取最后10KB内容
            f.seek(0, os.SEEK_END)
            pos = f.tell()
            size = 10240  # 读取最后10KB
            pos = max(pos - size, 0)
            f.seek(pos)
            lines = f.readlines()
            
            # 从后向前检查
            for line in reversed(lines):
                if "###log### 总用时:" in line:
                    return True
        return False
    except Exception as e:
        print(f"检查日志文件时出错: {e}")
        return False


def all_finish(scenarios_dir, copy_path):
    for file_name in os.listdir(scenarios_dir):
        if file_name.endswith(".json"):
            base_name = os.path.splitext(file_name)[0]    # 获取文件名（去掉后缀）
            destination_planner_dir = os.path.join(copy_path, base_name)    # 定义目标路径

            if os.path.exists(destination_planner_dir):
                if not is_scenario_completed(destination_planner_dir):
                    return False
            else:
                return False
    return True


def process_scenario_file(file_name, scenarios_dir, source_planner_dir, copy_path, env_name, t_wait):
    # global _process_ids
    base_name = os.path.splitext(file_name)[0]    # 获取文件名（去掉后缀）
    destination_planner_dir = os.path.join(copy_path, base_name)    # 定义目标路径

    if os.path.exists(destination_planner_dir):
        if is_scenario_completed(destination_planner_dir):
            print(f"目录 {destination_planner_dir} 已存在且场景已完成，跳过处理")
            return
            
        print(f"目录 {destination_planner_dir} 已存在，尝试删除")
        main_py = os.path.join(destination_planner_dir, "planner", "__main__.py")
        main_print = os.path.join(destination_planner_dir, "main_print.log")
        # try:
        #     os.remove(main_py)
        #     print(f"主函数 {main_py}未被占用，已成功删除")
        # # 删除目录失败，可能是文件被占用，如果被占用，则跳过    
        # except Exception as e:
        #     print(f"删除主函数 {main_py} 失败: {e}, 可能已有程序在运行, 跳过")
        #     return
        if is_file_locked(main_print):
            print(f"{main_print} 被占用, 跳过")
            return
        else:
            print(f"{main_print} 未被占用")
        try:
            shutil.rmtree(destination_planner_dir)
            print(f"删除目录 {destination_planner_dir} 成功")
        except Exception as e:
            print(f"删除目录 {destination_planner_dir} 失败: {e}, 可能已有程序在运行, 跳过")
            return

    shutil.copytree(source_planner_dir, destination_planner_dir)
    os.remove(os.path.join(destination_planner_dir, "planner", os.path.basename(__file__))) # 删除muti_run.py
    print(f"已复制{source_planner_dir}到{destination_planner_dir}")

    destination_json_dir = os.path.join(destination_planner_dir, "inputs", "Scenarios")
    os.makedirs(destination_json_dir, exist_ok=True)
    delete_file_dir(destination_json_dir, wait_print=False)

    outputs_dir = os.path.join(destination_planner_dir, "outputs")
    os.makedirs(outputs_dir, exist_ok=True)
    delete_file_dir(outputs_dir, wait_print=False)

    file_log = os.path.join(destination_planner_dir, "main_print.log")
    os.makedirs(file_log, exist_ok=True) if not file_log else None
    delete_file_dir(file_log, wait_print=False)

    # 获取并删除onsite_images_saved文件夹的所有子文件夹
    onsite_images_saved_dir = os.path.join(destination_planner_dir, "onsite_images_saved")
    for item in os.listdir(onsite_images_saved_dir):
        item_path = os.path.join(onsite_images_saved_dir, item)
        # 只删除子文件夹，保留文件
        if os.path.isdir(item_path):
            shutil.rmtree(item_path)
            print(f"已删除{item_path}")

    source_json_dir = os.path.join(scenarios_dir, file_name)
    destination_json_file = os.path.join(destination_json_dir, file_name)
    shutil.copyfile(source_json_dir, destination_json_file)

    # 检查当前运行的进程数是否已达到CPU核心数限制
    running_processes = update_running_processes()

    while running_processes >= CPU_COUNT:
        print(f"已达到CPU核心数限制({CPU_COUNT}), 等待{t_wait}秒后检查其他进程是否完成后再启动新进程")
        time.sleep(t_wait)
        running_processes = update_running_processes()

    # 初始启动命令：对每个复制好的JSON场景文件，启动新的命令行窗口运行planner
    command = f'start cmd /k "conda activate {env_name} && cd {destination_planner_dir} && python .\\planner && exit"'
    # before_pids = [p.pid for p in psutil.process_iter() if any(name in p.name().lower() for name in process_names)]  # 获取当前系统中所有匹配进程名的进程ID
    subprocess.Popen(command, shell=True, creationflags=subprocess.CREATE_NO_WINDOW)
                    
    # after_pids = [p.pid for p in psutil.process_iter() if any(name in p.name().lower() for name in process_names)]
    # new_pids = [pid for pid in after_pids if pid not in before_pids]
    # _process_ids = list(set(_process_ids + new_pids))  # 使用set消除重复项
    # print(f"_process_ids: {_process_ids}")

    print(f"\n启动进程: {base_name}，当前运行进程数: {update_running_processes()}/{CPU_COUNT}")
    print(f"等待{t_wait/5}秒后检查资源占用情况\n")
    time.sleep(t_wait/5)
    while check_resource_usage():
        print(f"资源占用过高，等待{t_wait}秒后重新检查")
        time.sleep(t_wait)  # 如果资源占用过高，等待一段时间后重新检查



def signal_handler(signum, frame):
    """处理Ctrl+C信号"""
    print(f"接收到退出信号，正在关闭...")
    cleanup_processes()
    os._exit(0)


def cleanup_processes():
    global _process_ids
    for pid in _process_ids:
        try:
            os.system(f'taskkill /F /PID {pid} /T 2>nul')
            print(f"已清除进程: {pid}")
            _process_ids = []   # 清理后清空_process_ids
        except:
            print(f'清理PID: {pid}失败')
            pass


def process():
    global _process_ids

    for file_name in os.listdir(scenarios_dir):    # 遍历目标文件夹中的JSON文件
        if file_name.endswith(".json"):
            # before_pids = [p.pid for p in psutil.process_iter() if any(name in p.name().lower() for name in process_names)]  # 获取当前系统中所有匹配进程名的进程ID
            process_scenario_file(file_name, scenarios_dir, os.path.dirname(dir_current_file), copy_path, env_name, t_wait)

            # after_pids = [p.pid for p in psutil.process_iter() if any(name in p.name().lower() for name in process_names)]
            # new_pids = [pid for pid in after_pids if pid not in before_pids]
            # _process_ids = list(set(_process_ids + new_pids))  # 使用set消除重复项
            # print(f"_process_ids: {_process_ids}")



def main():
    # 检查是否存在第二赛道-B卷文件夹
    if not os.path.exists(second_track_b_dir):
        print(f"文件夹 {second_track_b_dir} 不存在")
        exit(1)

    os.makedirs(copy_path, exist_ok=True)
    delete_file_dir(copy_path, default_delete=False)

    # 检查copy_path中是否存在是否存在第二赛道-B卷文件夹
    second_track_b_copy_dir = os.path.join(copy_path, "第二赛道-B卷")
    if not os.path.exists(second_track_b_copy_dir):
        # 复制第二赛道-B卷文件夹到copy_path
        shutil.copytree(second_track_b_dir, second_track_b_copy_dir)
        print(f"已复制{second_track_b_dir}到 {second_track_b_copy_dir}")

    global env_name
    env_name = get_env_name()
    if not os.path.exists(scenarios_dir):
        print(f"文件夹 {scenarios_dir} 不存在")
        exit(1)

    while not all_finish(scenarios_dir, copy_path):
        process()
        time.sleep(t_wait/5)
        print(f"{t_wait/5}秒后检查场景是否完成")
        


if __name__ == "__main__":
    # signal.signal(signal.SIGINT, signal_handler)
    # signal.signal(signal.SIGTERM, signal_handler)
    main()


